# 要添加一个新单元，输入 '# %%'
# 要添加一个新的标记单元，输入 '# %% [markdown]'
# 我发现写到最后，还是学弟写的那种样子。好麻烦啊，感觉应该一天就能写完，但自己在硬拖。
import copy
import torch
from torch.utils import data
from trainers.fedbase.FedClientBase import FedClientBase

import logging
from trainers.fedbase.FedClientBase import FedClientBase

# logging.basicConfig(level = logging.INFO,format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logging.basicConfig(format ='%(message)s')
logger = logging.getLogger()
logger.setLevel(logging.INFO)

class FedScaffoldClient(FedClientBase):
    """训练相关的参数本来应该是服务器发送过来的，包括模型、参数啥的。训练数据应该是本地的。
    这里为了方便，只模拟联邦学习的效果和逻辑，不模拟联邦学习的真是过程。
    1. 数据由服务器分成独立同分布，或者非独立同分布。然后分配给各个客户端。而非客户端在本地读取。
    2. 训练的模型、参数在客户端配置。
    3. 客户端完成训练过程、服务端完成聚合过程。客户端和服务端只发送和传输模型和优化器的权重参数。完成聚合和分配的通信逻辑。
    
    接收参数，完成训练。主要是一次训练的参数。
    每一个算法都应该提供梯度的返回值。即在客户端计算出梯度。其实在服务器计算也无可厚非。

    Args:
        FedClientBase (): 
    """    
    def __init__(self,datapair,model):
        """创建一个客户端对象。初始化客户端的本地参数

        Args:
            datapair (tuple(ndarray,ndarray)): input,output包括输入输出数据
            model (torch.model): 与客户端一致的模型
        """       
        super(FedScaffoldClient, self).__init__(datapair,model)    
 
        # 客户端基本参数，不需要配置--------------------------------------------------------------------------------------------
        # 服务器传过来的控制参数，和本地生成的控制参数。这个参数直接传聚合就行，不需要求变化。
        self.server_controls =  [torch.zeros_like(param) for param in self.model.parameters()]
        self.client_controls =  [torch.zeros_like(param) for param in self.model.parameters()]
        # 本地控制参数的变化
        print("FedScaffoldClient init----")

    def set_parameters(self,client_params):
        """从服务器加载训练的参数。并根据batch_size等参数划分数据集。

        Args:
            client_params (dict): 客户端进行训练的相关参数
        """        
        super().set_parameters(client_params)
        
        # scaffold配置参数-------------------------------------------------------------------------------------------------

        print("FedScaffoldClient set_parameters----")


    # 加载控制参数
    def load_controls(self,controls_list):
        self.server_controls = copy.deepcopy(controls_list)

    # 返回控制参数
    def get_controls(self):
        return copy.deepcopy(self.client_controls)
    
    # 用来训练当前的本地模型。每次训练返回损失值running_loss和running_accuracy
    def train(self):
        self.model.train()
        loss_sum = 0
        loss_times = 0.
        correct_sum = 0.
        total_size = 0
        # 保留一份原始参数
        old_param_list = [torch.clone(parm).detach() for parm in self.model.parameters()]
        # 计算两个不同的差异。用来衡量两个分布的差异程度。
        # 分布的差异越大，这个修正的服务应该越小，防止出现，过大修正，导致最后的发现。
        # 当分布的差异接近与1 的时候，进行全额的修正。因为太大的修正会导致最后本地模型中。该修正锁占的影响过大，最后出现发散。
        # 排除因为调整因子，起到过大的作用，导致每次更新波动太大，没办法收敛。
        server_client_controls_variance = self.get_params_norm(self.server_controls,self.client_controls)
        adaption = 1
        if server_client_controls_variance < 2 :
            adaption=1
        else:
            adaption = server_client_controls_variance

        
        for epoch in range(self.epochs):
            for batch_x,batch_y in self.train_dl:
                # 正向传播，计算准确度
                outputs = self.model(batch_x)
                # 计算loss
                loss = self.loss_func(outputs,batch_y.long())
                # 反向传播
                self.optimizer.zero_grad()
                loss.backward()
                # 梯度下降，修改模型梯度。将当前的控制参数放到梯度当中之后在修改。防止破坏现有的优化器，继续使用原来的方法进行梯度下降。
                for param, c, ci in zip(self.model.parameters(), self.server_controls, self.client_controls):
                    if param is None:
                        continue
                    # 修改现有的参数的grad。这相当与添加一个更新防线，与服务器上的更新方向一致。
                    # 我终于知道weight_control的意思了。他妈的。应该是用来平衡这里的更新次数和学习率。
                    # 使得经过K次更新后，正好累加了一次服务器上参数的更新方向。
                    param.grad.add_((c - ci)/adaption)
                self.optimizer.step()

                _,predicted = torch.max(outputs,1)
                correct_sum += (predicted == batch_y).sum().item()
                total_size += len(batch_y)
                loss_sum +=loss.item()
                loss_times+=1

        
        # 计算self.delta_model模型变化
        new_param_list = [torch.clone(parm).detach() for parm in self.model.parameters()]
        delta_model = []
        for new_param, old_param in zip(new_param_list, old_param_list):
            delta_model.append(new_param - old_param)
        
        old_controls = copy.deepcopy(self.client_controls)
        
        # 得到本地的控制参数。这个控制参数，表示，与服务器的控制参数的差，减去，变化的差。表示变化的剧烈程度
        # 相当于对本地的变化起到一定的抑制作用。待会执行以下，看看这个东西能不能，真的进行赋值。
        # 计算new_controls控制参数。我觉得这个有点多余。并不多余，非常有用
        weight_control = 1 /(self.inner_lr* loss_times)
        # 更新delta_controls和本地client_controls
        strategy = 1
        # 每一次本地更新的方向。如果将一下指放到grad中，相当于每次添加一个向服务器的方向更新的方向。
        if strategy==1:
            for i in range(len(self.client_controls)):
                self.client_controls[i] = -delta_model[i] * weight_control
        else:
            # 这玩意表示模型负变化的方向。也就是本地梯度的方向。但又考虑到了历史的梯度。
            # 在梯度下降的时候add_(-ci)表示减去上次的梯度值，抵消掉上次的梯度变化。这，好像又缺乏意义。
            # 是不是应该给个参数，减小一下影响更好。加一个参数0.2，用来微调 ，而不是真实影响。
            # 让历史的影响降低一些。之前的方向保留0.2.最新的方向保留0.8。
            # 在范数保持一致。事实上，原来的范数就是一致的。
            # 另外考虑到loss溢出的问题，估计多半是因为本地更新方向和服务器更新方向相差太多导致的。
            # 可以通过计算两个control的方法，显示这种差异。
            # 猜测出现这种情况的原因是，因为两个方向完全不一致，并且由于batch过多导致，更新次数太多，加大了这种差异。
            # 最终根本没办法收敛。所以，大胆猜测，可以减小本地更新次数。
            for i in range(len(self.client_controls)):
                self.client_controls[i] =-(self.server_controls[i]-self.client_controls[i] + delta_model[i] * weight_control)
                
        
        # 计算模型的共享两
        model_variance = self.get_params_norm(old_param_list,new_param_list)
        control_variance = self.get_params_norm(old_controls,self.client_controls)



        return loss_sum/loss_times,100*correct_sum/total_size,server_client_controls_variance,model_variance,control_variance
   
    # 用来测试训练的结果返回准确率
    def test(self):
        self.model.eval()
        correct = 0
        total =0
        with torch.no_grad():
            for features,labels in self.train_dl:
                outputs = self.model(features)
                # print(outputs.shape)
                _,predicted = torch.max(outputs,1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

        accuracy = 100*correct/total
        return accuracy